import os
from langchain_community.llms import Tongyi
from langsmith import Client
#
from langchain_core.tracers.context import tracing_v2_enabled
from langchain_core.prompts import ChatPromptTemplate

# 设置环境变量
os.environ["LANGCHAIN_TRACING_V2"] = "true"  # 启用V2追踪
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com" # LangSmith API端点
os.environ["LANGCHAIN_API_KEY"] = "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"  # 替换为你的LangSmith API密钥
"""上面的这些环境变量告诉LangChain SDK：
    要启用追踪功能
    数据发送到哪里
    如何认证身份
前面设置了这些环境变量后，当 langchain-core 被 import 时，会立即读取这些变量并全局启用追踪器。
之后任何 langchain 内部的运行（包括 chain.invoke、llm.invoke）都会自动把事件序列化并异步发送到 https://api.smith.langchain.com
"""
os.environ["LANGCHAIN_PROJECT"] = "pr-sunny-adrenalin-22"  # 设置项目名称，这个也是LangSmith上的配置

os.environ["DASHSCOPE_API_KEY"] = "sk-4e88cf4db3e14894bafaff606d296610"  # 替换为你的阿里云DashScope API密钥

def qwen_with_langsmith_example():
    """
    使用阿里千问模型并集成LangSmith进行追踪的示例
    """
    # 创建千问模型实例
    llm = Tongyi(
        model_name="qwen-turbo",  # 使用qwen-turbo，响应更快且成本更低
        temperature=0.7,
        dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY")
    )

    # 创建提示模板
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", "你是一个有帮助的AI助手，请用中文回答用户问题。回答要简洁明了。"),
        ("human", "{question}")
    ])

    # 创建链
    chain = prompt_template | llm

    # 启用LangSmith追踪
    with tracing_v2_enabled() as run_collector:
        # 在这里的代码会被自动追踪
        try:
            # 执行查询并追踪
            print("正在调用千问模型...")
            response = chain.invoke({
                "question": "请简要解释一下机器学习中的Transformer架构"
            })

            print("\n=== 模型回答 ===")
            print(response)

            # 获取运行ID的替代方法
            print("\n追踪信息已发送到LangSmith")
            print("请在 https://smith.langchain.com 查看详细追踪信息")

        except Exception as e:
            print(f"发生错误: {e}")

def qwen_with_explicit_tracing():
    """
    使用显式run_id追踪的示例
    """
    from langsmith import run_trees

    # 创建千问模型实例
    llm = Tongyi(
        model_name="qwen-turbo",
        temperature=0.7,
        dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY")
    )

    # 创建根运行
    root_run = run_trees.RunTree(
        name="Qwen-Demo-Run",
        run_type="chain",
        inputs={"question": "解释深度学习的基本概念"}
    )

    try:
        # 开始运行
        root_run.post()

        # 执行模型调用
        response = llm.invoke("请用通俗易懂的语言解释深度学习")

        # 记录输出
        root_run.end(outputs={"response": response})
        root_run.patch()

        print("\n=== 显式追踪示例回答 ===")
        print(response)
        print(f"\n运行ID: {root_run.id}")

        # 记录反馈
        client = Client()
        client.create_feedback(
            root_run.id,
            key="user_rating",
            score=4.5,
            comment="回答清晰易懂"
        )
        print("已记录反馈到LangSmith")

    except Exception as e:
        root_run.end(error=str(e))
        root_run.patch()
        print(f"发生错误: {e}")

def simple_qwen_example():
    """
    简化版示例，不依赖LangSmith
    """
    try:
        from langchain_community.llms import Tongyi

        llm = Tongyi(
            model_name="qwen-turbo",
            temperature=0.7,
            dashscope_api_key=os.environ.get("DASHSCOPE_API_KEY")
        )

        response = llm.invoke("请用一句话解释人工智能")
        print("\n=== 简化版示例回答 ===")
        print(response)

    except Exception as e:
        print(f"简化版示例出错: {e}")

if __name__ == "__main__":
    # 检查必要的环境变量
    required_vars = ["DASHSCOPE_API_KEY"]
    missing_vars = [var for var in required_vars if not os.getenv(var)]

    if missing_vars:
        print(f"请先设置以下环境变量: {', '.join(missing_vars)}")
        print("获取方式:")
        print("- DASHSCOPE_API_KEY: 访问 https://dashscope.aliyun.com 注册获取")
        print("- LANGCHAIN_API_KEY: 访问 https://smith.langchain.com 注册获取（可选）")
        print("\n正在运行简化版示例（不依赖LangSmith）...")
        simple_qwen_example()
    else:
        # 检查LangSmith配置
        if not os.getenv("LANGCHAIN_API_KEY"):
            print("警告: 未设置LANGCHAIN_API_KEY，将运行简化版示例")
            simple_qwen_example()
        else:
            print("=== 运行LangSmith集成示例 ===")
            qwen_with_langsmith_example()

            print("\n" + "="*50)
            print("=== 运行显式追踪示例 ===")
            qwen_with_explicit_tracing()